import torch
import torchvision
from torch import nn
from torch.nn import Conv2d, MaxPool2d, Flatten, Linear, Sequential
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

'''
基于PyTorch的卷积神经网络（CNN）模型训练流程，具体功能如下：
1. 加载CIFAR-10数据集并创建数据加载器。
2. 定义一个包含卷积层、池化层、展平层和全连接层的CNN模型。
3. 使用交叉熵损失函数计算模型输出与目标标签的损失值。
4. 通过反向传播计算梯度，并打印损失值。
'''
dataset = torchvision.datasets.CIFAR10("../../data", train=False,
                                       transform=torchvision.transforms.ToTensor(), download=True)
dataLoader = DataLoader(dataset,batch_size=64)

class Module(nn.Module):
    def __init__(self):
        super().__init__()
        self.module1 = Sequential(
            Conv2d(in_channels=3,out_channels=32,kernel_size=5,stride=1,padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=False),
            Conv2d(in_channels=32,out_channels=32,kernel_size=5,stride=1,padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=False),
            Conv2d(in_channels=32,out_channels=64,kernel_size=5,stride=1,padding=2),
            MaxPool2d(kernel_size=2, ceil_mode=False),
            Flatten(),
            Linear(1024,64),
            Linear(64,10)
        )

    def forward(self,x):
        x = self.module1(x)
        return x

module = Module()
print(module)
loss_cross = nn.CrossEntropyLoss()
for data in dataLoader:
    imgs,targets = data
    output = module(imgs)
    loss = loss_cross(output,targets)
    loss.backward()
    print(loss)